#include <test/testCases.cuh>
#include <assert.h>
#include <cmath>
#include <random>
#include <cstdio>
#include <iostream>
#include <utils/Profiler.h>
#include <Data/CSRMatrixOperator.cuh>
#include <utils/DataConvert.cuh>
#include <Smoother/JacobiSmoother.cuh>
#include <Precond/CSRPrecondAMG.cuh>

#define CHECK(call)                                                               \
{                                                                                 \
const cudaError_t error = call;                                               \
if (error != cudaSuccess)                                                     \
{                                                                             \
printf("Error: %s:%d, ", __FILE__, __LINE__);                             \
printf("code: %d, reason: %s\n", error, cudaGetErrorString(error));       \
exit(1);                                                                  \
}                                                                             \
}

void testSpmvHost() {
    std::vector<double> data{1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0};
    std::vector<double> xData{3.0, 2.0, 1.0};
    std::vector<double> expectedData{3.0, 4.0, 3.0};
    myamg::CSRMatrix<int, double> matrix(3, 3, 3);
    myamg::Vector<double> x(3);
    myamg::Vector<double> expected(3);
    myamg::Vector<double> result(3);
    matrix.loadFromDense(data, 3, 3);
    x.setValues(xData);
    expected.setValues(expectedData);
    myamg::matMulVec(matrix, x, result);
    for (auto i = 0; i < 3; i++) {
        assert(result.values()[i] == expected.values()[i]);
    }
}

void testSpmvDevice() {
    std::vector<double> data{1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 0.0, 0.0, 3.0};
    std::vector<double> xData{3.0, 2.0, 1.0};
    std::vector<double> expected{3.0, 4.0, 3.0};
    myamg::CSRMatrix<int, double, myamg::MatrixDeviceAllocator> matrix(3, 3, 3);
    myamg::Vector<double, myamg::MatrixDeviceAllocator> x(3);
    myamg::Vector<double, myamg::MatrixDeviceAllocator> result(3);
    matrix.loadFromDense(data, 3, 3);
    x.setValues(xData);
    myamg::matMulVec(matrix, x, result);
    auto hResult = new double[3];
    cudaMemcpy(hResult, result.values(), sizeof(double) * 3, cudaMemcpyDeviceToHost);
    for (auto i = 0; i < 3; i++) {
        assert(hResult[i] == expected.data()[i]);
    }
}

void profileSpmvHost() {
    const size_t iter = 100;
    std::string filename("Velx-0.csr");
    auto matrix = myamg::CSR2Matrix<int, double>(filename);
    auto n = matrix.numRows();
    auto Ap = matrix.rowPtr();
    auto Ai = matrix.colIndex();
    auto Av = matrix.values();
    std::vector<double> x(n);
    std::vector<double> y(n);
    constexpr size_t seed = 12345;
    std::mt19937_64 generator(seed);
    std::uniform_real_distribution<double> distribution(0.0, 1.0);
    for (size_t i = 0; i < n; i++) {
        x[i] = distribution(generator);
    }
    for (auto i = 0; i < iter; i++) {
        PROFILE_SCOPE("SPMV Host");
        myamg::spmvHost(n, Ap, Ai, Av, x.data(), y.data());
    }
}

void profileSpmvDevice() {
    const size_t iter = 100;
    std::string filename("Velx-0.csr");
    auto matrix = myamg::CSR2Matrix<int, double, myamg::MatrixDeviceAllocator>(filename);
    auto n = matrix.numRows();
    auto Ap = matrix.rowPtr();
    auto Ai = matrix.colIndex();
    auto Av = matrix.values();
    std::vector<double> x(n);
    constexpr size_t seed = 12345;
    std::mt19937_64 generator(seed);
    std::uniform_real_distribution<double> distribution(0.0, 1.0);
    for (size_t i = 0; i < n; i++) {
        x[i] = distribution(generator);
    }
    double *DeviceX;
    double *DeviceY;
    CHECK(cudaMalloc(&DeviceX, n * sizeof(double)));
    CHECK(cudaMalloc(&DeviceY, n * sizeof(double)));
    CHECK(cudaMemcpy(DeviceX, x.data(), n * sizeof(double), cudaMemcpyHostToDevice));
    dim3 block(512);
    dim3 grid((n + block.x - 1) / block.x);
    for (auto i = 0; i < iter; i++) {
        PROFILE_SCOPE("SPMV Device");
        myamg::spmvDevice<<<grid, block>>>(n, Ap, Ai, Av, DeviceX, DeviceY);
        CHECK(cudaDeviceSynchronize());
    }
}

void testYHMatrix2Matrix() {
    std::string filename("Velx-0.csr");
    myamg::CsrMatrixReader<int, double> reader(filename);
    auto yhMatrix = myamg::CSR2YHAMGCsrMatrix(filename);
    auto hostMatrix = myamg::YHMatrix2Matrix(*yhMatrix);
    auto deviceMatrix = myamg::YHMatrix2Matrix<myamg::MatrixDeviceAllocator>(*yhMatrix);
    auto nRows = hostMatrix.numRows();
    auto nElements = hostMatrix.numNonzeros();
    assert(nRows == reader.getNDim());
    assert(nElements == reader.getNData());
    for (auto i = 0; i < nRows + 1; i++) {
        assert(hostMatrix.rowPtr()[i] == yhMatrix->rowptr[i]);
    }
    for (auto i = 0; i < nElements; i++) {
        assert(hostMatrix.values()[i] == yhMatrix->values[i]);
    }
    for (auto i = 0; i < nElements; i++) {
        assert(hostMatrix.colIndex()[i] == hostMatrix.colIndex()[i]);
    }
    auto Ap = new int[nRows + 1];
    auto Ai = new int[nElements];
    auto Av = new double[nElements];
    CHECK(cudaMemcpy(Ap, deviceMatrix.rowPtr(), (nRows+1)*sizeof(int), cudaMemcpyDeviceToHost))
    CHECK(cudaMemcpy(Ai, deviceMatrix.colIndex(), nElements*sizeof(int), cudaMemcpyDeviceToHost))
    CHECK(cudaMemcpy(Av, deviceMatrix.values(), nElements*sizeof(double), cudaMemcpyDeviceToHost))
    CHECK(cudaDeviceSynchronize())
    for (auto i = 0; i < nRows + 1; i++) {
        assert(Ap[i] == yhMatrix->rowptr[i]);
    }
    for (auto i = 0; i < nElements; i++) {
        assert(Ai[i] == yhMatrix->colind[i]);
        assert(Av[i] == yhMatrix->values[i]);
    }
    delete[] Ap;
    delete[] Ai;
    delete[] Av;
}

__global__ void initZZero(double *z, size_t n) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        z[idx] = 0.0;
    }
}

__global__ void forwardUpdateX(double *b, double *z, int *Lp, int *Li, double *Lv, double *x, double *Drcpv, size_t n,
                               double omega) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        double temp = b[idx] + z[idx];
        for (int i = Lp[idx]; i < Lp[idx + 1]; i++) {
            temp -= Lv[i] * x[Li[i]];
        }
        x[idx] = temp * omega * Drcpv[idx];
    }
}

__global__ void forwardCalcResid(double *D, double *x, double *z, int *Ep, int *Ei, double *Ev, int *Up, int *Ui,
                                 double *Uv, double *r, size_t n, double theta) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        double temp = theta * D[idx] * x[idx] - z[idx];
        for (int i = Ep[idx]; i < Ep[idx + 1]; i++) {
            temp -= Ev[i] * x[Ei[i]];
        }
        for (int i = Up[idx]; i < Up[idx + 1]; i++) {
            temp -= Uv[i] * x[Ui[i]];
        }
        r[idx] = temp;
    }
}

__global__ void checkNan(double* data, const size_t n, bool* hasNan) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (hasNan) {
        return;
    }
    if (idx < n) {
        if (isnan(data[idx])) {
            hasNan[idx] = true;
        }
    }
}

bool checkNanDevice(double* data, const size_t n) {
    bool hasNan = false;
    dim3 block(512);
    dim3 grid((n + block.x - 1) / block.x);
    checkNan<<<grid, block>>>(data, n, &hasNan);
    CHECK(cudaDeviceSynchronize())
    return hasNan;
}

bool checkNanHost(double* data, const size_t n) {
    bool hasNan = false;
    for (auto i = 0; i < n; i++) {
        if (isnan(data[i])) {
            hasNan = true;
        }
    }
    return hasNan;
}

template<typename T>
void calcError(const T* A, const T* B, T& max, T& min, T& avg, const size_t n, const bool debug = false) {
    auto err = new T[n];
    for (auto i = 0; i < n; i++) {
        err[i] = std::abs(A[i] - B[i]);
    }
    T* errOriginal = nullptr;
    if (debug) {
        errOriginal = new T[n];
        std::copy(err, err + n, errOriginal);
    }
    std::sort(err, err + n);
    max = err[n - 1];
    min = err[0];
    if (debug) {
        auto minRow = 0, maxRow = 0;
        for (auto i = 0; i < n; i++) {
            if (errOriginal[i] == max) {
                maxRow = i;
            }
            if (errOriginal[i] == min) {
                minRow = i;
            }
        }
        std::cout << "max: " << maxRow << ", min: " << minRow << std::endl;
        delete[] errOriginal;
    }
    avg = 0;
    for (auto i = 0; i < n; i++) {
        avg += err[i];
    }
    avg = avg / n;
    delete[] err;
}

__global__ void printKernel(double* data, const size_t n) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        printf("%f\n", data[idx]);
    }
}

void profileSORDevice(const int batch, std::vector<std::unordered_map<std::string, double>>& error) {
    auto D = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/D_" + std::to_string(batch) + ".mtx");
    auto Drcpv = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/Drcpv_" + std::to_string(batch) + ".mtx");
    auto L = myamg::Mtx2Matrix<int, double, myamg::MatrixDeviceAllocator>("Catch/Mat/L_" + std::to_string(batch) + ".mtx");
    auto U = myamg::Mtx2Matrix<int, double, myamg::MatrixDeviceAllocator>("Catch/Mat/U_" + std::to_string(batch) + ".mtx");
    auto E = myamg::Mtx2Matrix<int, double, myamg::MatrixDeviceAllocator>("Catch/Mat/E_" + std::to_string(batch) + ".mtx");
    auto b = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/b_" + std::to_string(batch) + ".mtx");
    auto r = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/r_" + std::to_string(batch) + ".mtx");
    auto z = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/z_" + std::to_string(batch) + ".mtx");
    auto x = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/x_" + std::to_string(batch) + ".mtx");
    auto n = D.size();
    double omega = 1.0;
    double theta = 0.0;
    std::string gridSize = std::to_string(n) + "x" + std::to_string(n);

    dim3 block(512);
    dim3 grid((n + block.x - 1) / block.x);

    {
        YHAMG::Profiler::ScopedTimer wholeTimer(YHAMG::Profiler::instance(), "SOR Device" + gridSize, __FILE__);

        {
            YHAMG::Profiler::ScopedTimer timer(YHAMG::Profiler::instance(), "SOR Device init z zero" + gridSize, __FILE__);
            initZZero<<<grid, block>>>(z.values(), n);
            CHECK(cudaDeviceSynchronize());
        }
        {
            YHAMG::Profiler::ScopedTimer timer(YHAMG::Profiler::instance(), "SOR Device forward and resid" + gridSize, __FILE__);
            forwardUpdateX<<<grid, block>>>(b.values(), z.values(), L.rowPtr(), L.colIndex(), L.values(), x.values(),
                                            Drcpv.values(), n,
                                            omega);
            forwardCalcResid<<<grid, block>>>(D.values(), x.values(), z.values(), E.rowPtr(), E.colIndex(), E.values(),
                                              U.rowPtr(), U.colIndex(), U.values(), r.values(), n, theta);
            CHECK(cudaDeviceSynchronize());
        }
    }
    auto xResultHost = myamg::Mtx2Vector<double>("Catch/Vec/x_result_" + std::to_string(batch) + ".mtx");
    auto xResultDevice = new double[n];
    CHECK(cudaMemcpy(xResultDevice, x.values(), n * sizeof(double), cudaMemcpyDeviceToHost))
    CHECK(cudaDeviceSynchronize());
    double max = 0, min = 0, avg = 0;
    calcError(xResultHost.values(), xResultDevice, max, min, avg, n, false);
    auto& errorMap = error[batch-1];
    errorMap.emplace(std::make_pair("max", max));
    errorMap.emplace(std::make_pair("min", min));
    errorMap.emplace(std::make_pair("avg", avg));
}

void profileSOR() {
    constexpr auto batch = 70;
    std::vector<std::unordered_map<std::string, double>> error(batch);
    for (auto i = 1; i < batch+1; i++) {
        profileSORDevice(i, error);
    }
    double max = 0, min = 0, avg = 0;
    for (auto &x : error) {
        max += x["max"];
        min += x["min"];
        avg += x["avg"];
        std::cout << "max: " << x["max"] << ", min: " << x["min"] << ", avg: " << x["avg"] << std::endl;
    }
    avg /= batch;
    max /= batch;
    min /= batch;
    std::cout << "max: " << max << ", min: " << min << ", avg: " << avg << std::endl;
}

__global__ void forwardUpdateXandZ(double *b, double *z, int *Lp, int *Li, double *Lv, double *x, double *Drcpv, size_t n,
                               double omega) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        double temp = b[idx] + z[idx];
        for (int i = Lp[idx]; i < Lp[idx + 1]; i++) {
            temp -= Lv[i] * x[Li[i]];
        }
        x[idx] = temp * omega * Drcpv[idx];
        z[idx] = temp * (2 - omega) - z[idx];
    }
}

__global__ void forwardUpdateZExt(double* z, int* Ep, int* Ei, double* Ev, double* x, size_t n) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        double temp = z[idx];
        for (int i = Ep[idx]; i < Ep[idx + 1]; i++) {
            temp -= Ev[i] * x[Ei[i]];
        }
        z[idx] = temp;
    }
}

void calculateXError(int batch, double* x) {
    auto xResult = myamg::Mtx2Vector<double>("Catch/Vec/x_result_" + std::to_string(batch) + ".mtx");
    auto n = xResult.size();
    auto xResultDevice = new double[n];
    CHECK(cudaMemcpy(xResultDevice, x, n * sizeof(double), cudaMemcpyDeviceToHost))
    CHECK(cudaDeviceSynchronize())
    double max = 0, min = 0, avg = 0;
    calcError(xResult.values(), xResultDevice, max, min, avg, n, false);
    std::cout << "max: " << max << ", min: " << min << ", avg: " << avg << std::endl;
}

void searchSORDevice(int batch, std::vector<std::unordered_map<std::string, double>>& error, int maxIter) {
    auto D = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/D_" + std::to_string(batch) + ".mtx");
    auto Drcpv = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/Drcpv_" + std::to_string(batch) + ".mtx");
    auto L = myamg::Mtx2Matrix<int, double, myamg::MatrixDeviceAllocator>("Catch/Mat/L_" + std::to_string(batch) + ".mtx");
    auto U = myamg::Mtx2Matrix<int, double, myamg::MatrixDeviceAllocator>("Catch/Mat/U_" + std::to_string(batch) + ".mtx");
    auto E = myamg::Mtx2Matrix<int, double, myamg::MatrixDeviceAllocator>("Catch/Mat/E_" + std::to_string(batch) + ".mtx");
    auto b = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/b_" + std::to_string(batch) + ".mtx");
    auto r = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/r_" + std::to_string(batch) + ".mtx");
    auto z = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/z_" + std::to_string(batch) + ".mtx");
    auto x = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Catch/Vec/x_" + std::to_string(batch) + ".mtx");
    auto n = D.size();
    double omega = 1.0;
    double theta = 0.0;
    std::string gridSize = std::to_string(n) + "x" + std::to_string(n);

    dim3 block(512);
    dim3 grid((n + block.x - 1) / block.x);

    {
        {
            // 初始化z
            initZZero<<<grid, block>>>(z.values(), n);
            CHECK(cudaDeviceSynchronize());
        }
        {
            // SOR主循环
            int k = 0;
            while (true) {
                // 更新x，计算残差
                if (++k >= maxIter) {
                    forwardUpdateX<<<grid, block>>>(b.values(), z.values(), L.rowPtr(), L.colIndex(), L.values(), x.values(),
                                Drcpv.values(), n,
                                omega);
                    forwardCalcResid<<<grid, block>>>(D.values(), x.values(), z.values(), E.rowPtr(), E.colIndex(), E.values(),
                                                      U.rowPtr(), U.colIndex(), U.values(), r.values(), n, theta);
                    CHECK(cudaDeviceSynchronize());
                    calculateXError(batch, x.values());
                    break;
                } else {
                    // 更新x和z
                    forwardUpdateXandZ<<<grid, block>>>(b.values(), z.values(), L.rowPtr(), L.colIndex(), L.values(), x.values(),
                                Drcpv.values(), n,
                                omega);
                    forwardUpdateZExt<<<grid, block>>>(z.values(), E.rowPtr(), E.colIndex(), E.values(), x.values(), n);
                    CHECK(cudaDeviceSynchronize());
                    calculateXError(batch, x.values());
                }
            }
        }
    }
}

int testYHAMGJacobi(int argc, char* argv[])
{
    std::string inputFile("Vely-1.csr");

    std::cout << "Loading input file: " << inputFile << std::endl;

    MPI_Init(&argc, &argv);

    // AMG Param
    int max_levels = 20;
    int coarse_size = 8;
    double strength_threshold = 0.25;
    int aggressive_levels = 50;
    int coarsen_type = 0; // 0: HMIS, 1: PMIS
    int interp_type = 0; // 0: Long-Range Interpolation, 1: Smoothed Aggregation
    int interp_min_elements = 4;
    int interp_max_elements = 4;
    double truncation_factor = 0.1;
    double sparsification_threshold = 0.01;
    int cycle_type = 0; // 0: V-Cycle, 1: W-Cycle, 2: F-Cycle
    int pre_sweeps = 1;
    int post_sweeps = 1;
    int coarse_sweeps = 2;
    int smooth_type = 0; // 0: Jacobi, 1: SOR, 2: ILU, 3: Chebyshev

    double jacobi_factor = 0.75;

    int sor_type = 2; // 0: Forward, 1: Backword, 2: Symmetric
    double sor_factor = 1.0;

    int ilu_maxfil = 0;
    double ilu_droptol = 0.01;

    int chebyshev_order = 4;
    double chebyshev_eigen_ratio = 0.3;

    int gmres_restart = 10;
    int max_iters = 500;
    double tolerance = 1.0e-08;

    bool print_stats = 1;

    MPI_Comm comm = MPI_COMM_WORLD;

    auto readA = myamg::CSR2YHAMGCsrMatrix(inputFile);
    auto readB = myamg::CSRRhs2YHAMGVector(inputFile);

    auto A = std::make_unique<YHAMG::ParCSRMatrix>(*readA);
    auto b = std::make_unique<YHAMG::ParVector>(*readB);
    auto x = std::make_unique<YHAMG::ParVector>(comm);
    x->Resize(b->local.size);
    x->Fill(0.0);

    auto P = std::make_unique<YHAMG::ParCSRPrecondAMG>();

    P->MaxLevels = max_levels;
    P->CoarseSize = coarse_size;
    P->StrengthThreshold = strength_threshold;
    P->AggressiveLevels = aggressive_levels;
    P->CoarsenType = coarsen_type;
    P->InterpType = interp_type;
    P->InterpMinElements = interp_min_elements;
    P->InterpMaxElements = interp_max_elements;
    P->TruncationFactor = truncation_factor;
    P->SparsificationThreshold = sparsification_threshold;
    P->CycleType = cycle_type;
    P->PreSweeps = pre_sweeps; // pre步数
    P->PostSweeps = post_sweeps; // post步数
    P->CoarseSweeps = coarse_sweeps; // 最后一层步数
    P->SmoothType = smooth_type;
    P->JacobiFactor = jacobi_factor;
    P->SORType = sor_type;
    P->SORFactor = sor_factor;
    P->ILUMaxFillins = ilu_maxfil;
    P->ILUDropTolerance = ilu_droptol;
    P->ChebyshevEigenRatio = chebyshev_eigen_ratio;
    P->PrintStats = print_stats;

    std::cout << "\n# AMG Setup Phase\n";
    std::cout << "--------------------------------------------------\n";

    P->Setup(*A);

    std::cout << "\n# BiCGStab Solve Phase\n";
    std::cout << "--------------------------------------------------\n";

    auto solverCG = std::make_unique<YHAMG::ParSolverBCGS>(max_iters, tolerance, print_stats);

    int iter;
    double relres;

    (*solverCG)(*A, *P, *b, *x, iter, relres);

    MPI_Finalize();
    return 0;
}

__global__ void jacobiUpdateZ(double* b, int* Ap, int* Ai, double* Av, double* x, double* z, size_t n) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        double temp = b[idx];
        for (int i = Ap[idx]; i < Ap[idx + 1]; i++) {
            temp -= Av[i] * x[Ai[i]];
        }
        z[idx] = temp;
    }
}

__global__ void jacobiUpdateX(double* x, double* Drcpv, double* z, size_t n, double omega) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        x[idx] = x[idx] + omega * Drcpv[idx] * z[idx];
    }
}

__global__ void jacobiUpdateXZero(double* x, double* Drcpv, double* b, size_t n, double omega) {
    const size_t idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        x[idx] = omega * Drcpv[idx] * b[idx];
    }
}

void testJacobiDevice(const int batch, std::vector<std::unordered_map<std::string, double>>& error) {
    auto A = myamg::Mtx2Matrix<int, double, myamg::MatrixDeviceAllocator>("Jacobi/Mat/A_" + std::to_string(batch) + ".mtx");
    auto Drcpv = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Jacobi/Vec/Drcpv_" + std::to_string(batch) + ".mtx");
    auto x = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Jacobi/Vec/x_" + std::to_string(batch) + ".mtx");
    auto b = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Jacobi/Vec/b_" + std::to_string(batch) + ".mtx");
    auto z = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Jacobi/Vec/z_" + std::to_string(batch) + ".mtx");
    auto x_result = myamg::Mtx2Vector<double>("Jacobi/Vec/x_result_" + std::to_string(batch) + ".mtx");
    auto param = myamg::Mtx2Vector<double>("Jacobi/Vec/param_" + std::to_string(batch) + ".mtx");
    double omega = param.values()[0];
    int step = static_cast<int>(param.values()[1]);
    bool x0zero = static_cast<int>(param.values()[2]);
    // std::cout << "omega: " << omega << ", step: " << step << ", x0zero: " << x0zero << std::endl;

    auto n = A.numRows();
    auto nNonzero = A.numNonzeros();

    dim3 block(512);
    dim3 grid((n + block.x - 1) / block.x);

    auto gridSize = std::to_string(n) + "x" + std::to_string(n);

    {
        YHAMG::Profiler::ScopedTimer wholeTimer(YHAMG::Profiler::instance(), "Jacobi Device", "");
        for (auto k = 0; k < step; k++) {
            if (!x0zero || k > 0) {
                {
                    YHAMG::Profiler::ScopedTimer timer(YHAMG::Profiler::instance(), "Jacobi Update z Device", "");
                    // z = b - A * x
                    jacobiUpdateZ<<<grid, block>>>(b.values(), A.rowPtr(), A.colIndex(), A.values(), x.values(), z.values(), n);
                    CHECK(cudaDeviceSynchronize());
                }
                {
                    YHAMG::Profiler::ScopedTimer timer(YHAMG::Profiler::instance(), "Jacobi Update x Device", "");
                    // x = x + omega * Drcpv * z
                    jacobiUpdateX<<<grid, block>>>(x.values(), Drcpv.values(), z.values(), n, omega);
                    CHECK(cudaDeviceSynchronize());
                }
            } else {
                {
                    YHAMG::Profiler::ScopedTimer timer(YHAMG::Profiler::instance(), "Jacobi Update x Device (x0 = 0)", "");
                    // x = omega * Drcpv * b
                    jacobiUpdateXZero<<<grid, block>>>(x.values(), Drcpv.values(), b.values(), n, omega);
                    CHECK(cudaDeviceSynchronize());
                }
            }
        }
    }

    auto xResultDevice = new double[n];
    CHECK(cudaMemcpy(xResultDevice, x.values(), n * sizeof(double), cudaMemcpyDeviceToHost))
    CHECK(cudaDeviceSynchronize());
    double max = 0, min = 0, avg = 0;
    calcError(x_result.values(), xResultDevice, max, min, avg, n, false);
    auto& errorMap = error[batch-1];
    errorMap.emplace(std::make_pair("max", max));
    errorMap.emplace(std::make_pair("min", min));
    errorMap.emplace(std::make_pair("avg", avg));
    delete[] xResultDevice;
}

void testJacobiDeviceAll(int batch) {
    std::vector<std::unordered_map<std::string, double>> error(batch);
    for (auto i = 1; i < batch+1; i++) {
        testJacobiDevice(i, error);
    }
    double max = 0, min = 0, avg = 0;
    for (auto &x : error) {
        max += x["max"];
        min += x["min"];
        avg += x["avg"];
        std::cout << "max: " << x["max"] << ", min: " << x["min"] << ", avg: " << x["avg"] << std::endl;
    }
    avg /= batch;
    max /= batch;
    min /= batch;
    std::cout << "max: " << max << ", min: " << min << ", avg: " << avg << std::endl;
}

void testJacobiSmoother() {
    auto A = myamg::Mtx2Matrix<int, double, myamg::MatrixDeviceAllocator>("Jacobi/Mat/A_1.mtx");
    auto b = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Jacobi/Vec/b_1.mtx");
    auto x = myamg::Mtx2Vector<double, myamg::MatrixDeviceAllocator>("Jacobi/Vec/x_1.mtx");
    auto x_result = myamg::Mtx2Vector<double>("Jacobi/Vec/x_result_1.mtx");
    auto param = myamg::Mtx2Vector<double>("Jacobi/Vec/param_1.mtx");
    double omega = param.values()[0];
    int step = static_cast<int>(param.values()[1]);
    bool x0zero = static_cast<int>(param.values()[2]);
    myamg::JacobiSmoother<int, double, myamg::MatrixDeviceAllocator> smoother{omega};
    smoother.setup(A);
    smoother.pre(A, b, x, step, x0zero);
    CHECK(cudaDeviceSynchronize());
    auto xResultDevice = new double[A.numRows()];
    CHECK(cudaMemcpy(xResultDevice, x.values(), A.numRows() * sizeof(double), cudaMemcpyDeviceToHost))
    CHECK(cudaDeviceSynchronize());
    double max = 0, min = 0, avg = 0;
    calcError(x_result.values(), xResultDevice, max, min, avg, A.numRows(), false);
    constexpr double threshold = 1e-9;
    assert(max < threshold);
    assert(min < threshold);
    assert(avg < threshold);
}

void testVectorClear() {
    myamg::Vector<double> vectorHostDouble(10);
    myamg::Vector<int> vectorHostInt(10);
    myamg::Vector<float> vectorHostFloat(10);
    myamg::Vector<double, myamg::MatrixDeviceAllocator> vectorDeviceDouble(10);
    myamg::Vector<int, myamg::MatrixDeviceAllocator> vectorDeviceInt(10);
    myamg::Vector<float, myamg::MatrixDeviceAllocator> vectorDeviceFloat(10);

    std::vector<double> dataDouble(10, 1.0);
    std::vector<int> dataInt(10, 1);
    std::vector<float> dataFloat(10, 1.0f);
    vectorHostDouble.setValues(dataDouble);
    vectorHostInt.setValues(dataInt);
    vectorHostFloat.setValues(dataFloat);
    vectorDeviceDouble.setValues(dataDouble);
    vectorDeviceInt.setValues(dataInt);
    vectorDeviceFloat.setValues(dataFloat);

    std::vector<double> deviceDataDouble(10);
    std::vector<int> deviceDataInt(10);
    std::vector<float> deviceDataFloat(10);

    CHECK(cudaMemcpy(deviceDataDouble.data(), vectorDeviceDouble.values(), 10 * sizeof(double), cudaMemcpyDeviceToHost));
    CHECK(cudaMemcpy(deviceDataInt.data(), vectorDeviceInt.values(), 10 * sizeof(int), cudaMemcpyDeviceToHost));
    CHECK(cudaMemcpy(deviceDataFloat.data(), vectorDeviceFloat.values(), 10 * sizeof(float), cudaMemcpyDeviceToHost));
    CHECK(cudaDeviceSynchronize());

    for (int i = 0; i < 10; i++) {
        assert(vectorHostDouble.values()[i] != 0);
        assert(vectorHostInt.values()[i] != 0);
        assert(vectorHostFloat.values()[i] != 0);
        assert(deviceDataDouble[i] != 0);
        assert(deviceDataInt[i] != 0);
        assert(deviceDataFloat[i] != 0);
    }

    myamg::vectorClear(vectorHostDouble);
    myamg::vectorClear(vectorHostInt);
    myamg::vectorClear(vectorHostFloat);
    myamg::vectorClear(vectorDeviceDouble);
    myamg::vectorClear(vectorDeviceInt);
    myamg::vectorClear(vectorDeviceFloat);

    CHECK(cudaMemcpy(deviceDataDouble.data(), vectorDeviceDouble.values(), 10 * sizeof(double), cudaMemcpyDeviceToHost));
    CHECK(cudaMemcpy(deviceDataInt.data(), vectorDeviceInt.values(), 10 * sizeof(int), cudaMemcpyDeviceToHost));
    CHECK(cudaMemcpy(deviceDataFloat.data(), vectorDeviceFloat.values(), 10 * sizeof(float), cudaMemcpyDeviceToHost));
    CHECK(cudaDeviceSynchronize());

    for (int i = 0; i < 10; i++) {
        assert(vectorHostDouble.values()[i] == 0);
        assert(vectorHostInt.values()[i] == 0);
        assert(vectorHostFloat.values()[i] == 0);
        assert(deviceDataDouble[i] == 0);
        assert(deviceDataInt[i] == 0);
        assert(deviceDataFloat[i] == 0);
    }
}

void testRAP() {
    std::cout << "Testing RAP function..." << std::endl;

    using Integer = int;
    using Real = double;
    using Allocator = myamg::MatrixDeviceAllocator;
    using MatrixType = myamg::CSRMatrix<Integer, Real, Allocator>;

    const Real tol = 1e-10;

    {
        Integer A_nnz = 4;
        Integer *A_rp = new Integer[3]{0, 2, 4};
        Integer *A_ci = new Integer[4]{0, 1, 0, 1};
        Real *A_v = new Real[4]{2.0, -1.0, -1.0, 2.0};
        MatrixType A(2, 2, A_nnz);
        A.setRowPtr(A_rp);
        A.setColIndex(A_ci);
        A.setValues(A_v);

        Integer P_nnz = 2;
        Integer *P_rp = new Integer[3]{0, 1, 2};
        Integer *P_ci = new Integer[2]{0, 0};
        Real *P_v = new Real[2]{1.0, 1.0};
        MatrixType P(2, 1, P_nnz);
        P.setRowPtr(P_rp);
        P.setColIndex(P_ci);
        P.setValues(P_v);

        Integer R_nnz = 2;
        Integer *R_rp = new Integer[2]{0, 2};
        Integer *R_ci = new Integer[2]{0, 1};
        Real *R_v = new Real[2]{1.0, 1.0};
        MatrixType R(1, 2, R_nnz);
        R.setRowPtr(R_rp);
        R.setColIndex(R_ci);
        R.setValues(R_v);

        MatrixType C;
        myamg::RugeStuben<Integer, Real, Allocator>::RAP(R, A, P, C);

        assert(C.numRows() == 1);
        assert(C.numCols() == 1);
        assert(C.numNonzeros() == 1);

        Integer *C_rp_host = new Integer[C.numRows() + 1];
        Integer *C_ci_host = new Integer[C.numNonzeros()];
        Real *C_v_host = new Real[C.numNonzeros()];

        cudaMemcpy(C_rp_host, C.rowPtr(), (C.numRows() + 1) * sizeof(Integer), cudaMemcpyDeviceToHost);
        cudaMemcpy(C_ci_host, C.colIndex(), C.numNonzeros() * sizeof(Integer), cudaMemcpyDeviceToHost);
        cudaMemcpy(C_v_host, C.values(), C.numNonzeros() * sizeof(Real), cudaMemcpyDeviceToHost);

        assert(C_rp_host[0] == 0);
        assert(C_rp_host[1] == 1);
        assert(C_ci_host[0] == 0);
        assert(std::abs(C_v_host[0] - 2.0) < tol);

        delete[] C_rp_host;
        delete[] C_ci_host;
        delete[] C_v_host;
        delete[] A_rp; delete[] A_ci; delete[] A_v;
        delete[] P_rp; delete[] P_ci; delete[] P_v;
        delete[] R_rp; delete[] R_ci; delete[] R_v;

        std::cout << "  Test 1 (2x2 simple) PASSED" << std::endl;
    }

    {
        Integer A_nnz = 12;
        Integer *A_rp = new Integer[5]{0, 3, 6, 9, 12};
        Integer *A_ci = new Integer[12]{0, 1, 3,
                                        0, 1, 2,
                                        1, 2, 3,
                                        0, 2, 3};
        Real *A_v = new Real[12]{4.0, -1.0, -1.0,
                                 -1.0, 4.0, -1.0,
                                 -1.0, 4.0, -1.0,
                                 -1.0, -1.0, 4.0};
        MatrixType A(4, 4, A_nnz);
        A.setRowPtr(A_rp);
        A.setColIndex(A_ci);
        A.setValues(A_v);

        Integer P_nnz = 6;
        Integer *P_rp = new Integer[5]{0, 1, 3, 5, 6};
        Integer *P_ci = new Integer[6]{0, 0, 1, 0, 1, 1};
        Real *P_v = new Real[6]{1.0, 0.5, 0.5, 0.5, 0.5, 1.0};
        MatrixType P(4, 2, P_nnz);
        P.setRowPtr(P_rp);
        P.setColIndex(P_ci);
        P.setValues(P_v);

        Integer R_nnz = 6;
        Integer *R_rp = new Integer[3]{0, 3, 6};
        Integer *R_ci = new Integer[6]{0, 1, 2, 1, 2, 3};
        Real *R_v = new Real[6]{1.0, 0.5, 0.5, 0.5, 0.5, 1.0};
        MatrixType R(2, 4, R_nnz);
        R.setRowPtr(R_rp);
        R.setColIndex(R_ci);
        R.setValues(R_v);

        MatrixType C;
        myamg::RugeStuben<Integer, Real, Allocator>::RAP(R, A, P, C);

        assert(C.numRows() == 2);
        assert(C.numCols() == 2);
        assert(C.numNonzeros() == 4);

        Integer *C_rp_host = new Integer[C.numRows() + 1];
        Integer *C_ci_host = new Integer[C.numNonzeros()];
        Real *C_v_host = new Real[C.numNonzeros()];

        cudaMemcpy(C_rp_host, C.rowPtr(), (C.numRows() + 1) * sizeof(Integer), cudaMemcpyDeviceToHost);
        cudaMemcpy(C_ci_host, C.colIndex(), C.numNonzeros() * sizeof(Integer), cudaMemcpyDeviceToHost);
        cudaMemcpy(C_v_host, C.values(), C.numNonzeros() * sizeof(Real), cudaMemcpyDeviceToHost);

        assert(C_rp_host[0] == 0);
        assert(C_rp_host[1] == 2);
        assert(C_rp_host[2] == 4);

        assert(C_ci_host[0] == 0);
        assert(std::abs(C_v_host[0] - 4.5) < tol);
        assert(C_ci_host[1] == 1);
        assert(std::abs(C_v_host[1] - (-0.5)) < tol);

        assert(C_ci_host[2] == 0);
        assert(std::abs(C_v_host[2] - (-0.5)) < tol);
        assert(C_ci_host[3] == 1);
        assert(std::abs(C_v_host[3] - 4.5) < tol);

        delete[] C_rp_host;
        delete[] C_ci_host;
        delete[] C_v_host;
        delete[] A_rp; delete[] A_ci; delete[] A_v;
        delete[] P_rp; delete[] P_ci; delete[] P_v;
        delete[] R_rp; delete[] R_ci; delete[] R_v;

        std::cout << "  Test 2 (4x4 to 2x2) PASSED" << std::endl;
    }

    {
        Integer A_nnz = 5;
        Integer *A_rp = new Integer[4]{0, 2, 3, 5};
        Integer *A_ci = new Integer[5]{0, 2, 1, 0, 2};
        Real *A_v = new Real[5]{3.0, -1.0, 2.0, -1.0, 3.0};
        MatrixType A(3, 3, A_nnz);
        A.setRowPtr(A_rp);
        A.setColIndex(A_ci);
        A.setValues(A_v);

        Integer P_nnz = 3;
        Integer *P_rp = new Integer[4]{0, 1, 2, 3};
        Integer *P_ci = new Integer[3]{0, 0, 0};
        Real *P_v = new Real[3]{1.0, 0.5, 1.0};
        MatrixType P(3, 1, P_nnz);
        P.setRowPtr(P_rp);
        P.setColIndex(P_ci);
        P.setValues(P_v);

        Integer R_nnz = 3;
        Integer *R_rp = new Integer[2]{0, 3};
        Integer *R_ci = new Integer[3]{0, 1, 2};
        Real *R_v = new Real[3]{1.0, 0.5, 1.0};
        MatrixType R(1, 3, R_nnz);
        R.setRowPtr(R_rp);
        R.setColIndex(R_ci);
        R.setValues(R_v);

        MatrixType C;
        myamg::RugeStuben<Integer, Real, Allocator>::RAP(R, A, P, C);

        assert(C.numRows() == 1);
        assert(C.numCols() == 1);
        assert(C.numNonzeros() == 1);

        Real *C_v_host = new Real[C.numNonzeros()];
        cudaMemcpy(C_v_host, C.values(), C.numNonzeros() * sizeof(Real), cudaMemcpyDeviceToHost);

        assert(std::abs(C_v_host[0] - 4.5) < tol);

        delete[] C_v_host;
        delete[] A_rp; delete[] A_ci; delete[] A_v;
        delete[] P_rp; delete[] P_ci; delete[] P_v;
        delete[] R_rp; delete[] R_ci; delete[] R_v;

        std::cout << "  Test 3 (3x3 with zeros) PASSED" << std::endl;
    }

    std::cout << "All RAP tests PASSED!" << std::endl;
}

void profileDot(const int size, const int iter) {
    myamg::Vector<double> xHost(size);
    myamg::Vector<double> yHost(size);
    myamg::Vector<double, myamg::MatrixDeviceAllocator> xDevice(size);
    myamg::Vector<double, myamg::MatrixDeviceAllocator> yDevice(size);
    std::vector<double> xData(size);
    std::vector<double> yData(size);
    for (auto i = 0; i < size; i++) {
        xData[i] = static_cast<double>(i + 1);
    }
    for (auto i = size-1; i >= 0 ; i--) {
        yData[i] = static_cast<double>(i + 1);
    }
    double hostResult = 0;
    double deviceResult = 0;
    for (auto i = 0; i < iter; i++) {
        PROFILE_SCOPE("Host Dot");
        hostResult = myamg::vectorDotProduct(xHost, yHost);
    }
    for (auto i = 0; i < iter; i++) {
        PROFILE_SCOPE("Device Dot");
        deviceResult = myamg::vectorDotProduct(xDevice, yDevice);
    }
    assert(std::abs(hostResult - deviceResult) < 1e-6);
}

int profileJacobiBiCGStab(int argc, char* argv[])
{
    std::string inputFile("Vely-1.csr");

    std::cout << "Loading input file: " << inputFile << std::endl;

    // AMG Param
    int max_levels = 20;
    int coarse_size = 8;
    double strength_threshold = 0.25;
    int aggressive_levels = 50;
    int coarsen_type = 0; // 0: HMIS, 1: PMIS
    int interp_type = 0; // 0: Long-Range Interpolation, 1: Smoothed Aggregation
    int interp_min_elements = 4;
    int interp_max_elements = 4;
    double truncation_factor = 0.1;
    double sparsification_threshold = 0.01;
    int cycle_type = 0; // 0: V-Cycle, 1: W-Cycle, 2: F-Cycle
    int pre_sweeps = 1;
    int post_sweeps = 1;
    int coarse_sweeps = 2;
    int smooth_type = 0; // 0: Jacobi, 1: SOR, 2: ILU, 3: Chebyshev

    double jacobi_factor = 0.75;

    int sor_type = 2; // 0: Forward, 1: Backword, 2: Symmetric
    double sor_factor = 1.0;

    int ilu_maxfil = 0;
    double ilu_droptol = 0.01;

    int chebyshev_order = 4;
    double chebyshev_eigen_ratio = 0.3;

    int gmres_restart = 10;
    int max_iters = 500;
    double tolerance = 1.0e-08;

    bool print_stats = 1;

    MPI_Comm comm = MPI_COMM_WORLD;

    auto readA = myamg::CSR2YHAMGCsrMatrix(inputFile);
    auto readB = myamg::CSRRhs2YHAMGVector(inputFile);

    auto A = std::make_unique<YHAMG::ParCSRMatrix>(*readA);
    auto b = std::make_unique<YHAMG::ParVector>(*readB);
    auto x = std::make_unique<YHAMG::ParVector>(comm);
    x->Resize(b->local.size);
    x->Fill(0.0);

    auto P = std::make_unique<YHAMG::ParCSRPrecondAMG>();

    P->MaxLevels = max_levels;
    P->CoarseSize = coarse_size;
    P->StrengthThreshold = strength_threshold;
    P->AggressiveLevels = aggressive_levels;
    P->CoarsenType = coarsen_type;
    P->InterpType = interp_type;
    P->InterpMinElements = interp_min_elements;
    P->InterpMaxElements = interp_max_elements;
    P->TruncationFactor = truncation_factor;
    P->SparsificationThreshold = sparsification_threshold;
    P->CycleType = cycle_type;
    P->PreSweeps = pre_sweeps; // pre步数
    P->PostSweeps = post_sweeps; // post步数
    P->CoarseSweeps = coarse_sweeps; // 最后一层步数
    P->SmoothType = smooth_type;
    P->JacobiFactor = jacobi_factor;
    P->SORType = sor_type;
    P->SORFactor = sor_factor;
    P->ILUMaxFillins = ilu_maxfil;
    P->ILUDropTolerance = ilu_droptol;
    P->ChebyshevEigenRatio = chebyshev_eigen_ratio;
    P->PrintStats = print_stats;

    std::cout << "\n# AMG Setup Phase\n";
    std::cout << "--------------------------------------------------\n";

    {
        // PROFILE_SCOPE("Jacobi-AMG Setup CPU");
        P->Setup(*A);
    }

    std::cout << "\n# BiCGStab Solve Phase\n";
    std::cout << "--------------------------------------------------\n";

    auto solverCG = std::make_unique<YHAMG::ParSolverBCGS>(max_iters, tolerance, print_stats);

    int iter;
    double relres;

    {
        // PROFILE_SCOPE("Jacobi-AMG-BiCGStab Solve CPU");
        (*solverCG)(*A, *P, *b, *x, iter, relres);
    }


    return 0;
}

int profileSORBiCGStab(int argc, char* argv[])
{
    std::string inputFile("Vely-1.csr");

    std::cout << "Loading input file: " << inputFile << std::endl;



    // AMG Param
    int max_levels = 20;
    int coarse_size = 8;
    double strength_threshold = 0.25;
    int aggressive_levels = 50;
    int coarsen_type = 0; // 0: HMIS, 1: PMIS
    int interp_type = 0; // 0: Long-Range Interpolation, 1: Smoothed Aggregation
    int interp_min_elements = 4;
    int interp_max_elements = 4;
    double truncation_factor = 0.1;
    double sparsification_threshold = 0.01;
    int cycle_type = 0; // 0: V-Cycle, 1: W-Cycle, 2: F-Cycle
    int pre_sweeps = 1;
    int post_sweeps = 1;
    int coarse_sweeps = 2;
    int smooth_type = 1; // 0: Jacobi, 1: SOR, 2: ILU, 3: Chebyshev

    double jacobi_factor = 0.75;

    int sor_type = 2; // 0: Forward, 1: Backword, 2: Symmetric
    double sor_factor = 1.0;

    int ilu_maxfil = 0;
    double ilu_droptol = 0.01;

    int chebyshev_order = 4;
    double chebyshev_eigen_ratio = 0.3;

    int gmres_restart = 10;
    int max_iters = 500;
    double tolerance = 1.0e-08;

    bool print_stats = 1;

    MPI_Comm comm = MPI_COMM_WORLD;

    auto readA = myamg::CSR2YHAMGCsrMatrix(inputFile);
    auto readB = myamg::CSRRhs2YHAMGVector(inputFile);

    auto A = std::make_unique<YHAMG::ParCSRMatrix>(*readA);
    auto b = std::make_unique<YHAMG::ParVector>(*readB);
    auto x = std::make_unique<YHAMG::ParVector>(comm);
    x->Resize(b->local.size);
    x->Fill(0.0);

    auto P = std::make_unique<YHAMG::ParCSRPrecondAMG>();

    P->MaxLevels = max_levels;
    P->CoarseSize = coarse_size;
    P->StrengthThreshold = strength_threshold;
    P->AggressiveLevels = aggressive_levels;
    P->CoarsenType = coarsen_type;
    P->InterpType = interp_type;
    P->InterpMinElements = interp_min_elements;
    P->InterpMaxElements = interp_max_elements;
    P->TruncationFactor = truncation_factor;
    P->SparsificationThreshold = sparsification_threshold;
    P->CycleType = cycle_type;
    P->PreSweeps = pre_sweeps; // pre步数
    P->PostSweeps = post_sweeps; // post步数
    P->CoarseSweeps = coarse_sweeps; // 最后一层步数
    P->SmoothType = smooth_type;
    P->JacobiFactor = jacobi_factor;
    P->SORType = sor_type;
    P->SORFactor = sor_factor;
    P->ILUMaxFillins = ilu_maxfil;
    P->ILUDropTolerance = ilu_droptol;
    P->ChebyshevEigenRatio = chebyshev_eigen_ratio;
    P->PrintStats = print_stats;

    std::cout << "\n# AMG Setup Phase\n";
    std::cout << "--------------------------------------------------\n";

    {
        // PROFILE_SCOPE("SOR-AMG Setup CPU");
        P->Setup(*A);
    }

    std::cout << "\n# BiCGStab Solve Phase\n";
    std::cout << "--------------------------------------------------\n";

    auto solverCG = std::make_unique<YHAMG::ParSolverBCGS>(max_iters, tolerance, print_stats);

    int iter;
    double relres;

    {
        // PROFILE_SCOPE("SOR-AMG-BiCGStab Solve CPU");
        (*solverCG)(*A, *P, *b, *x, iter, relres);
    }


    return 0;
}